package org.leiyang.common.LoadBalancer;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.LoadBalancerRequest;
import org.springframework.cloud.client.loadbalancer.Request;
import org.springframework.cloud.client.loadbalancer.Response;
import org.springframework.cloud.client.loadbalancer.reactive.ReactiveLoadBalancer;
import org.springframework.cloud.loadbalancer.blocking.client.BlockingLoadBalancerClient;
import reactor.core.publisher.Mono;

import java.io.IOException;
import java.net.URI;
import java.util.Objects;
import java.util.concurrent.CompletableFuture;

public class CustomBlockingLoadBalancerClient extends BlockingLoadBalancerClient {
    private static final Logger log = LoggerFactory.getLogger(CustomBlockingLoadBalancerClient.class);
    private final ReactiveLoadBalancer.Factory<ServiceInstance> loadBalancerClientFactory;
    private final BlockingLoadBalancerClient blockingLoadBalancerClient;

    public CustomBlockingLoadBalancerClient(ReactiveLoadBalancer.Factory<ServiceInstance> loadBalancerClientFactory,
                                            BlockingLoadBalancerClient blockingLoadBalancerClient) {
        super(loadBalancerClientFactory);
        this.loadBalancerClientFactory = loadBalancerClientFactory;
        this.blockingLoadBalancerClient = blockingLoadBalancerClient;
    }

    @Override
    public <T> ServiceInstance choose(String serviceId, Request<T> request) {
        log.info("【mark】Execute CustomBlockingLoadBalancerClient....");
        ReactiveLoadBalancer<ServiceInstance> loadBalancer = loadBalancerClientFactory.getInstance(serviceId);
        if (loadBalancer == null) {
            return null;
        }
        //异步转同步
        CompletableFuture<Response<ServiceInstance>> f = CompletableFuture.supplyAsync(() -> Mono.from(loadBalancer.choose(request)).block());
        Response<ServiceInstance> loadBalancerResponse = null;

        try {
            loadBalancerResponse = f.get();
        } catch (Exception e) {
            log.error("Get ServiceInstance from CustomBlockingLoadBalancerClient Failed, ", e);
        }

        return Objects.isNull(loadBalancerResponse) ? (blockingLoadBalancerClient.choose(serviceId, request)) : loadBalancerResponse.getServer();
    }

    @Override
    public <T> T execute(String serviceId, LoadBalancerRequest<T> request) throws IOException {
        return super.execute(serviceId, request);
    }

    @Override
    public <T> T execute(String serviceId, ServiceInstance serviceInstance, LoadBalancerRequest<T> request) throws IOException {
        return super.execute(serviceId, serviceInstance, request);
    }

    @Override
    public URI reconstructURI(ServiceInstance serviceInstance, URI original) {
        return super.reconstructURI(serviceInstance, original);
    }

    @Override
    public ServiceInstance choose(String serviceId) {
        return super.choose(serviceId);
    }
}
